import os
import argparse

import torch

from modelscope import pipeline, Tasks
from modelscope.outputs import OutputKeys


def extract_mask(video_input_path='https://modelscope.oss-cn-beijing.aliyuncs.com/test/videos/video_matting_test.mp4',
                 output_path='matting_out.mp4', model='damo/cv_effnetv2_video-human-matting', device='gpu'):
    video_matting = pipeline(Tasks.video_human_matting, model=model, device=device)
    result_status = video_matting({'video_input_path': video_input_path, 'output_path': output_path})
    return result_status[OutputKeys.MASKS]

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--video', type=str, required=True, help='input video file')
    parser.add_argument('--result-type', type=str, default='matte', choices=['matte'],
                        help='matte - save the alpha matte; fg - save the foreground')
    parser.add_argument('--model-name', type=str, default='/vepfs_train2/meichaoyang/model/damo/cv_effnetv2_video-human-matting', help='model path')

    print('Get CMD Arguments...')
    args = parser.parse_args()

    if not os.path.exists(args.video):
        print('Cannot find the input video: {0}'.format(args.video))
        exit()

    if torch.cuda.device_count() > 0:
        device = 'cuda'
    else:
        device = 'cpu'

    result = os.path.join("data","output", os.path.splitext(os.path.basename(args.video))[0] + '_effnetv2_{0}.mp4'.format(args.result_type))
    extract_mask(args.video, result, args.model_name, device)
